# data/fundus_seg_dataset.py
import os
from typing import List, Tuple
import numpy as np
from PIL import Image

import torch
from torch.utils.data import Dataset
from torchvision.transforms import Resize

class fundusSeg(Dataset):
    """
    0-shot 推理数据集：
    - 直接从 image_paths 读取图片
    - 自动 resize 到 SAM encoder 输入尺寸
    - 使用整图框作为提示框
    - modal 固定为传入的 modal_id (Fundus=6)
    - organ_name 只能是 "OpticCup" 或 "OpticDisc"
    - 返回的 data 字段键保持与原项目一致：img / box / modal / organ / name
    """
    def __init__(
        self,
        image_paths: List[str],
        organ_name: str,            # "OpticCup" or "OpticDisc"
        modal_id: int = 6,          # Fundus
        sam_input_size: int = 256,  # 会由外部 SAM 模型读取后传入
    ):
        assert organ_name in ("OpticCup", "OpticDisc"), \
            f"organ_name must be OpticCup or OpticDisc, got {organ_name}"
        self.image_paths = image_paths
        self.organ_name = organ_name
        self.modal_id = modal_id
        self.S = sam_input_size
        self.resize = Resize((self.S, self.S), antialias=True)

        # 可选导入 taxonomy；若没有就用 fallback
        self._have_taxonomy = False
        try:
            from data.taxonomy import (
                organ_level_1_dict, organ_level_1_map,
                organ_level_2_dict, organ_level_2_map,
                organ_level_3_dict, organ_level_3_map,
                task_idx
            )
            self._have_taxonomy = True
            self._tax = (organ_level_1_dict, organ_level_1_map,
                         organ_level_2_dict, organ_level_2_map,
                         organ_level_3_dict, organ_level_3_map,
                         task_idx)
        except Exception:
            # 简单兜底（不崩溃即可）
            self._have_taxonomy = False

        # 预计算 organ 索引
        self.organ_tuple = self._make_organ_tuple(self.organ_name)

    def _make_organ_tuple(self, organ: str) -> Tuple[int, int, int, int]:
        """根据 taxonomy 生成 organ 层级索引；若无 taxonomy 则 fallback。"""
        organ_clean = organ.rstrip("0123456789")  # 与原逻辑对齐
        if not self._have_taxonomy:
            # fallback: 给个稳定可用的占位
            organ_level_1 = 0
            organ_level_2 = 0
            organ_level_3 = 0
            organ_level_4 = 0 if organ_clean.lower() == "opticcup" else 1
            return (organ_level_1, organ_level_2, organ_level_3, organ_level_4)

        (organ_level_1_dict, organ_level_1_map,
         organ_level_2_dict, organ_level_2_map,
         organ_level_3_dict, organ_level_3_map,
         task_idx) = self._tax

        # body parts
        for k, v in organ_level_1_dict.items():
            if organ_clean in v:
                organ_level_1 = organ_level_1_map[k]
                break
        # body subregions
        for k, v in organ_level_2_dict.items():
            if organ_clean in v:
                organ_level_2 = organ_level_2_map[k]
                break
        # organs and tissues
        for k, v in organ_level_3_dict.items():
            if organ_clean in v:
                organ_level_3 = organ_level_3_map[k]
                break
        organ_level_4 = task_idx[organ_clean]
        return (organ_level_1, organ_level_2, organ_level_3, organ_level_4)

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx: int):
        p = self.image_paths[idx]
        pil = Image.open(p).convert("RGB")
        img_np = np.array(pil)              # H,W,3  uint8
        img_np = img_np.transpose(2, 0, 1)  # 3,H,W
        img = torch.tensor(img_np).float()  # 与原项目一致：不做/255 归一化
        img = self.resize(img)              # 3,S,S

        # 整图框（按 resize 后的尺寸）
        S = img.shape[-1]
        box = torch.tensor([0, 0, S-1, S-1], dtype=torch.float32)

        # modal: 直接整数或张量都可（下游会转 tensor）
        modal = self.modal_id

        # organ: 按项目风格返回 4 元组索引
        organ_tuple = self.organ_tuple

        data = {
            "img": img,            # (3,S,S) float32
            "box": box,            # (4,)
            "modal": modal,        # int
            "organ": organ_tuple,  # (l1,l2,l3,task_idx)
            "name": p,             # 原始路径方便保存
        }

        # 0-shot 无 GT：返回个 dummy label（与评测签名兼容）
        dummy_gt = torch.zeros(1, S, S, dtype=torch.long)
        return data, dummy_gt
